import numpy as np
import quimb.tensor as qtn
from quimb.tensor import Tensor, TensorNetwork
from quimb.tensor.tensor_core import (
	rand_uuid
)

I_o2 = np.identity(4)
F_o3 = np.array([  # [F0, F1, F2, F3]
	[[1.0, 0, 0, 0],
	 [0, 0, 0, 0],
	 [0, 0, 0, 0],
	 [0, 0, 0, 1]],

	[[0, 0, 0, 1],
	 [0, 0, 0, 0],
	 [0, 0, 0, 0],
	 [1, 0, 0, 0]],

	[[0, 0, 0, 0],
	 [0, 1, 0, 0],
	 [0, 0, 1, 0],
	 [0, 0, 0, 0]],

	[[0, 0, 0, 0],
	 [0, 0, 1, 0],
	 [0, -1, 0, 0],
	 [0, 0, 0, 0]]
])
G_o3 = np.array([  # [G0, G1, G2, G3]
	[[1.0, 0, 0, 0],
	 [0, 1, 0, 0],
	 [0, 0, 0, 0],
	 [0, 0, 0, 0]],

	[[0, 0, 0, 0],
	 [0, 0, 0, 0],
	 [0, 0, 1, 0],
	 [0, 0, 0, 1]],

	[[0, 1, 0, 0],
	 [1, 0, 0, 0],
	 [0, 0, 0, 0],
	 [0, 0, 0, 0]],

	[[0, 0, 0, 0],
	 [0, 0, 0, 0],
	 [0, 0, 0, 1],
	 [0, 0, -1, 0]]
])
Rx_o2_func = lambda _theta: np.array([
	[1, 0, 0, 0],
	[0, 1, 0, 0],
	[0, 0., np.cos(_theta), -np.sin(_theta)],
	[0, 0., np.sin(_theta), np.cos(_theta)],
])
Rz_o2_func = lambda _theta: np.array([
	[1, 0, 0, 0],
	[0, np.cos(_theta), -np.sin(_theta), 0],
	[0, np.sin(_theta), np.cos(_theta), 0],
	[0, 0, 0, 1]
])

def MPO_1(
		_N_qubit: int,
		_theta: float,
		_upper_ind_id="k{}", _lower_ind_id="b{}", _site_tag_id="I{}",
		**tn_opts,
):
	_upper_inds, _lower_inds = map(_upper_ind_id.format, range(_N_qubit)), map(_lower_ind_id.format, range(_N_qubit)),
	_site_tags = list(map(_site_tag_id.format, range(_N_qubit)))
	_Rx_o2 = Rx_o2_func(_theta)

	def gen_tensors():
		for _i in range(_N_qubit):
			_upper_ind, _lower_ind = next(_upper_inds), next(_lower_inds)
			_k_ix, _site_tag = rand_uuid(), _site_tags[_i]
			yield Tensor(_Rx_o2, inds=(_k_ix, _lower_ind), tags=[_site_tag])
			if _i % 2: yield Tensor(G_o3, inds=(_m_ix, _upper_ind, _k_ix), tags=[_site_tag])
			else:
				if _i == _N_qubit - 1: yield Tensor(I_o2, inds=(_upper_ind, _k_ix), tags=[_site_tag])
				else:
					_m_ix = rand_uuid()
					yield Tensor(F_o3, inds=(_m_ix, _upper_ind, _k_ix), tags=[_site_tag])

	return TensorNetwork(gen_tensors(), virtual=True, **tn_opts)

def inv_MPO_1(
		_N_qubit: int,
		_theta: float,
		_upper_ind_id="k{}", _lower_ind_id="b{}", _site_tag_id="I{}",
		**tn_opts,
):
	_upper_inds, _lower_inds = map(_upper_ind_id.format, range(_N_qubit)), map(_lower_ind_id.format, range(_N_qubit)),
	_site_tags = list(map(_site_tag_id.format, range(_N_qubit)))
	_Rx_o2 = Rx_o2_func(-_theta)

	def gen_tensors():
		for _i in range(_N_qubit):
			_upper_ind, _lower_ind = next(_upper_inds), next(_lower_inds)
			_k_ix, _site_tag = rand_uuid(), _site_tags[_i]
			if _i % 2: yield Tensor(G_o3, inds=(_m_ix, _k_ix, _lower_ind), tags=[_site_tag])
			else:
				if _i == _N_qubit - 1: yield Tensor(I_o2, inds=(_k_ix, _lower_ind), tags=[_site_tag])
				else:
					_m_ix = rand_uuid()
					yield Tensor(F_o3, inds=(_m_ix, _k_ix, _lower_ind), tags=[_site_tag])
			yield Tensor(_Rx_o2, inds=(_upper_ind, _k_ix), tags=[_site_tag])

	return TensorNetwork(gen_tensors(), virtual=True, **tn_opts)

def MPO_2(
		_N_qubit: int,
		_theta: float,
		_upper_ind_id="k{}", _lower_ind_id="b{}", _site_tag_id="I{}",
		**tn_opts,
):
	_upper_inds, _lower_inds = map(_upper_ind_id.format, range(_N_qubit)), map(_lower_ind_id.format, range(_N_qubit)),
	_site_tags = list(map(_site_tag_id.format, range(_N_qubit)))
	_Rz_o2 = Rz_o2_func(_theta)

	def gen_tensors():
		for _i in range(_N_qubit):
			_upper_ind, _lower_ind, _site_tag = next(_upper_inds), next(_lower_inds), _site_tags[_i]
			if _i % 2:
				_k_ix = rand_uuid()
				yield Tensor(_Rz_o2, inds=(_k_ix, _lower_ind), tags=[_site_tag])
				yield Tensor(G_o3, inds=(_m_ix, _upper_ind, _k_ix), tags=[_site_tag])
			else:
				if _i == _N_qubit - 1: yield Tensor(I_o2, inds=(_upper_ind, _lower_ind), tags=[_site_tag])
				else:
					_m_ix = rand_uuid()
					yield Tensor(F_o3, inds=(_m_ix, _upper_ind, _lower_ind), tags=[_site_tag])

	return TensorNetwork(gen_tensors(), virtual=True, **tn_opts)

def inv_MPO_2(
		_N_qubit: int,
		_theta: float,
		_upper_ind_id="k{}", _lower_ind_id="b{}", _site_tag_id="I{}",
		**tn_opts,
):
	_upper_inds, _lower_inds = map(_upper_ind_id.format, range(_N_qubit)), map(_lower_ind_id.format, range(_N_qubit)),
	_site_tags = list(map(_site_tag_id.format, range(_N_qubit)))
	_Rz_o2 = Rz_o2_func(-_theta)

	def gen_tensors():
		for _i in range(_N_qubit):
			_upper_ind, _lower_ind, _site_tag = next(_upper_inds), next(_lower_inds), _site_tags[_i]
			if _i % 2:
				_k_ix = rand_uuid()
				yield Tensor(G_o3, inds=(_m_ix, _k_ix, _lower_ind), tags=[_site_tag])
				yield Tensor(_Rz_o2, inds=(_upper_ind, _k_ix), tags=[_site_tag])
			else:
				if _i == _N_qubit - 1: yield Tensor(I_o2, inds=(_upper_ind, _lower_ind), tags=[_site_tag])
				else:
					_m_ix = rand_uuid()
					yield Tensor(F_o3, inds=(_m_ix, _upper_ind, _lower_ind), tags=[_site_tag])

	return TensorNetwork(gen_tensors(), virtual=True, **tn_opts)

def MPO_3(
		_N_qubit: int,
		_upper_ind_id="k{}", _lower_ind_id="b{}", _site_tag_id="I{}",
		**tn_opts,
):
	_upper_inds, _lower_inds, _site_tags = map(_upper_ind_id.format, range(_N_qubit)), \
		map(_lower_ind_id.format, range(_N_qubit)), map(_site_tag_id.format, range(_N_qubit))

	def gen_tensors():
		for _i in range(_N_qubit):
			_upper_ind, _lower_ind, _site_tag = next(_upper_inds), next(_lower_inds), next(_site_tags)
			if _i % 2:
				if _i == _N_qubit - 1: yield Tensor(I_o2, inds=(_upper_ind, _lower_ind), tags=[_site_tag])
				else:
					_m_ix = rand_uuid()
					yield Tensor(F_o3, inds=(_m_ix, _upper_ind, _lower_ind), tags=[_site_tag])
			else:
				if _i == 0: yield Tensor(I_o2, inds=(_upper_ind, _lower_ind), tags=[_site_tag])
				else: yield Tensor(G_o3, inds=(_m_ix, _upper_ind, _lower_ind), tags=[_site_tag])

	return TensorNetwork(gen_tensors(), virtual=True, **tn_opts)

def MPO_4(
		_N_qubit: int,
		_theta: float,
		_upper_ind_id="k{}", _lower_ind_id="b{}", _site_tag_id="I{}",
		**tn_opts,
):
	_upper_inds, _lower_inds, _site_tags = map(_upper_ind_id.format, range(_N_qubit)), \
		map(_lower_ind_id.format, range(_N_qubit)), map(_site_tag_id.format, range(_N_qubit))
	_Rz_o2 = Rz_o2_func(_theta)

	def gen_tensors():
		for _i in range(_N_qubit):
			_upper_ind, _lower_ind, _site_tag = next(_upper_inds), next(_lower_inds), next(_site_tags)
			if _i % 2:
				if _i == _N_qubit - 1: yield Tensor(I_o2, inds=(_upper_ind, _lower_ind), tags=[_site_tag])
				else:
					_m_ix = rand_uuid()
					yield Tensor(F_o3, inds=(_m_ix, _upper_ind, _lower_ind), tags=[_site_tag])
			else:
				if _i == 0: yield Tensor(I_o2, inds=(_upper_ind, _lower_ind), tags=[_site_tag])
				else:
					_k_ix = rand_uuid()
					yield Tensor(_Rz_o2, inds=(_k_ix, _lower_ind), tags=[_site_tag])
					yield Tensor(G_o3, inds=(_m_ix, _upper_ind, _k_ix), tags=[_site_tag])

	return TensorNetwork(gen_tensors(), virtual=True, **tn_opts)

def inv_MPO_4(
		_N_qubit: int,
		_theta: float,
		_upper_ind_id="k{}", _lower_ind_id="b{}", _site_tag_id="I{}",
		**tn_opts,
):
	_upper_inds, _lower_inds, _site_tags = map(_upper_ind_id.format, range(_N_qubit)), \
		map(_lower_ind_id.format, range(_N_qubit)), map(_site_tag_id.format, range(_N_qubit))
	_Rz_o2 = Rz_o2_func(-_theta)

	def gen_tensors():
		for _i in range(_N_qubit):
			_upper_ind, _lower_ind, _site_tag = next(_upper_inds), next(_lower_inds), next(_site_tags)
			if _i % 2:
				if _i == _N_qubit - 1: yield Tensor(I_o2, inds=(_upper_ind, _lower_ind), tags=[_site_tag])
				else:
					_m_ix = rand_uuid()
					yield Tensor(F_o3, inds=(_m_ix, _upper_ind, _lower_ind), tags=[_site_tag])
			else:
				if _i == 0: yield Tensor(I_o2, inds=(_upper_ind, _lower_ind), tags=[_site_tag])
				else:
					_k_ix = rand_uuid()
					yield Tensor(G_o3, inds=(_m_ix, _k_ix, _lower_ind), tags=[_site_tag])
					yield Tensor(_Rz_o2, inds=(_upper_ind, _k_ix), tags=[_site_tag])

	return TensorNetwork(gen_tensors(), virtual=True, **tn_opts)

# 10×3的二维矩阵
layer1_lam1_mat = np.array([
	[0.0007188384538, 0.0026219497414, 0.0015496680743],  # qubit 1
	[0.0013697548210, 0.0015972844218, 0.0026786081192],  # qubit 2
	[0.0009840379484, 0.0015345232914, 0.002027448252],  # qubit 3
	[0.001063884090, 0.0001065144408, 0.0017631475265],  # qubit 4
	[0.0, 0.0, 0.0015495526495],  # qubit 5
	[0.0, 0.0023902359479, 0.0010958132238],  # qubit 6
	[0.0020257855184, 0.0021400855405, 0.0005788805915],  # qubit 7
	[0.0019446767565, 0.0014794883469, 0.0004834382590],  # qubit 8
	[0.0014522106481, 0.0009070535380, 0.0023264494057],  # qubit 9
	[0.0006858528793, 0.0010735326374, 0.0009854359356]  # qubit 10
])
layer2_lam1_mat = np.array([
	[0.0, 0.0014693600715, 0.0026627276453],
	[0.0009620591518, 0.0010564395207, 0.0009723465495],
	[0.0013388355933, 0.0010427760033, 0.0],
	[0.0023105048745, 0.0014919297884, 0.0002082390859],
	[0.0009261012701, 0.0007442025724, 0.0023443828473],
	[0.0, 0.0007993187098, 0.0017095711268],
	[0.0006089087327, 0.0003111400198, 0.0],
	[0.0021851855014, 0.0006081573813, 0.0010763953485],
	[0.0006056925257, 0.0008321612132, 0.0010019294656],
	[0.0, 0.0017389696194, 0.0026542233490]
])

# 9×9的二维矩阵
layer1_lam2_mat = np.array([
	# qubit 1-2
	[0.0004171448785, 0.0004000124513, 0.0, 0.0013544305447, 0.0006013215169, 0.0001360847576, 0.0002374506778, 0.0,
	 0.0],
	# qubit 2-3
	[0.0, 0.0014591789675, 0.0001143396436, 0.0003231157903, 0.0003047123740, 0.0009723767925, 0.0009163673808,
	 0.0014617357731, 0.0001052694463],
	# qubit 3-4
	[0.0007171705339, 0.0004325833089, 0.0017211740069, 0.0010317935262, 0.0007849610247, 0.0009968302758, 0.0,
	 0.0002774239964, 0.0],
	# qubit 4-5
	[0.0, 0.0004145140852, 0.0, 0.0002399826614, 0.0014391631248, 0.0, 0.0, 0.0, 0.0011609124234],
	# qubit 5-6
	[0.0, 0.0020571240593, 0.0, 0.0011758015780, 0.0003494956822, 0.0005244052712, 0.0003472416886, 0.0002517186965,
	 0.0008080793202],
	# qubit 6-7
	[0.0004431773843, 0.000341754743, 0.0002432753686, 0.0009004397971, 0.0, 0.0008233462545, 0.0007495748739,
	 0.0007152515147, 0.0012941373561],
	# qubit 7-8
	[0.0006800975765, 0.0011992236472, 0.0010657734054, 0.000569312825, 0.000110333258, 0.000585147777, 0.0008180286985,
	 0.0, 0.0002920344835],
	# qubit 8-9
	[0.0003609533666, 0.0006946695859, 0.0011090370933, 0.0004346218145, 0.0, 0.0001599411535, 0.0006980513828,
	 0.0, 0.0013271429858],
	# qubit 9-10
	[0.0009791413973, 0.0011589174283, 0.0001831596684, 0.0003006216905, 0.0004491016202, 0.0010136247483,
	 0.0010611984571, 0.0, 0.0012903240539]
])
layer2_lam2_mat = np.array([
	[0.0004237835913, 0.0003630373210, 0.0002848231897, 0.0, 0.0, 0.0005966276473, 0.0, 0.0007309942553,
	 0.0000015157442],
	[0.0015575541386, 0.0, 0.0005837913797, 0.0001372174639, 0.0001625124553, 0.0009019187150, 0.0, 0.0002304339894,
	 0.0000869601531],
	[0.0008090860326, 0.0009637438292, 0.0008743308693, 0.0002868450783, 0.0013018963990, 0.0000797314206,
	 0.0002635981122, 0.0001441857959, 0.0017425011303],
	[0.0007077367641, 0.0002408123508, 0.0006666894811, 0.0000454529163, 0.0004523831237, 0.0003603372457, 0.0,
	 0.0002380569629, 0.0004160620513],
	[0.0, 0.0011914765728, 0.0008437638617, 0.0001668267636, 0.0009173528423, 0.0, 0.0010055930809, 0.0000652947879,
	 0.0007740090482],
	[0.0003856405199, 0.0017265374157, 0.0011718305575, 0.0003206445689, 0.0003806542922, 0.0009424159528,
	 0.0005749526574, 0.0009080912917, 0.0012534980664],
	[0.0007253097676, 0.0, 0.0004465530507, 0.0006274975328, 0.0, 0.0006127668507, 0.0012447982973, 0.0006178629978,
	 0.0006609217767],
	[0.0002185124723, 0.0005693242109, 0.0007241280299, 0.0007282630750, 0.0010641609267, 0.0001106725915, 0.0,
	 0.0020606323308, 0.0003172806200],
	[0.0002306927820, 0.0008086945582, 0.0014672289588, 0.0, 0.0009904281035, 0.0002488933528, 0.0003356747248, 0.0,
	 0.0003655776749]
])

from SPLM import SPLM_ptm_mpo

def trotter_step_MPO(
		_N_qubit: int, _J: float, _h: float, _dt: float,
		_upper_ind_id="k{}", _lower_ind_id="b{}", _site_tag_id="I{}",
		_max_bond: int = 200, _cutoff: float = 1e-10, method='direct'
):
	_theta1, _theta2 = 2 * _h * _dt, -2 * _J * _dt
	_ts_tn = MPO_1(_N_qubit, _theta1,
				   _upper_ind_id="_ts_a{}", _lower_ind_id=_lower_ind_id, _site_tag_id=_site_tag_id)
	_ts_tn |= MPO_2(_N_qubit, _theta2,
					_upper_ind_id="_ts_b{}", _lower_ind_id="_ts_a{}", _site_tag_id=_site_tag_id)
	_ts_tn |= MPO_3(_N_qubit,
					_upper_ind_id="_ts_c{}", _lower_ind_id="_ts_b{}", _site_tag_id=_site_tag_id)
	_ts_tn |= MPO_4(_N_qubit, _theta2,
					_upper_ind_id=_upper_ind_id, _lower_ind_id="_ts_c{}", _site_tag_id=_site_tag_id)
	_site_tags = tuple(map(_site_tag_id.format, range(_N_qubit)))
	qtn.tensor_network_1d_compress(_ts_tn, max_bond=_max_bond, cutoff=_cutoff, site_tags=_site_tags, method=method,
								   inplace=True)

	return _ts_tn

def noisy_trotter_step_MPO(
		_N_qubit: int, _J: float, _h: float, _dt: float,
		_layer1_lam1_mat, _layer1_lam2_mat, _layer2_lam1_mat, _layer2_lam2_mat,
		_upper_ind_id="k{}", _lower_ind_id="b{}", _site_tag_id="I{}",
		_max_bond: int = 200, _cutoff: float = 1e-10, method='direct'
):
	_site_tags = tuple(map(_site_tag_id.format, range(_N_qubit)))
	_theta1, _theta2 = 2 * _h * _dt, -2 * _J * _dt
	_nts_tn = MPO_1(_N_qubit, _theta1,
					_upper_ind_id="_nts_a{}", _lower_ind_id=_lower_ind_id, _site_tag_id=_site_tag_id)
	_nts_tn |= SPLM_ptm_mpo(_N_qubit, _layer1_lam1_mat, _layer1_lam2_mat,
							upper_ind_id="_nts_b{}", lower_ind_id="_nts_a{}", site_tag_id=_site_tag_id)
	_nts_tn |= MPO_2(_N_qubit, _theta2,
					 _upper_ind_id="_nts_c{}", _lower_ind_id="_nts_b{}", _site_tag_id=_site_tag_id)
	_nts_tn |= SPLM_ptm_mpo(_N_qubit, _layer1_lam1_mat, _layer1_lam2_mat,
							upper_ind_id="_nts_d{}", lower_ind_id="_nts_c{}", site_tag_id=_site_tag_id)
	qtn.tensor_network_1d_compress(_nts_tn, max_bond=_max_bond, cutoff=_cutoff, site_tags=_site_tags, inplace=True,
								   method=method)
	_nts_tn |= MPO_3(_N_qubit,
					 _upper_ind_id="_nts_e{}", _lower_ind_id="_nts_d{}", _site_tag_id=_site_tag_id)
	_nts_tn |= SPLM_ptm_mpo(_N_qubit, _layer2_lam1_mat, _layer2_lam2_mat,
							upper_ind_id="_nts_f{}", lower_ind_id="_nts_e{}", site_tag_id=_site_tag_id)
	_nts_tn |= MPO_4(_N_qubit, _theta2,
					 _upper_ind_id="_nts_g{}", _lower_ind_id="_nts_f{}", _site_tag_id=_site_tag_id)
	_nts_tn |= SPLM_ptm_mpo(_N_qubit, _layer2_lam1_mat, _layer2_lam2_mat,
							upper_ind_id=_upper_ind_id, lower_ind_id="_nts_g{}", site_tag_id=_site_tag_id)
	qtn.tensor_network_1d_compress(_nts_tn, max_bond=_max_bond, cutoff=_cutoff, site_tags=_site_tags, method=method,
								   inplace=True)

	return _nts_tn

def TEM_trotter_steps(
		_N_qubit, _J: float, _h: float, _dt: float,
		_layer1_lam1_mat, _layer1_lam2_mat, _layer2_lam1_mat, _layer2_lam2_mat,
		_upper_ind_id="k{}", _lower_ind_id="b{}", _site_tag_id="I{}",
		_max_bond: int = 200, _cutoff: float = 1e-10, method='direct'
):
	from SPLM import inv_SPLM_ptm_mpo

	_theta1, _theta2 = 2 * _h * _dt, -2 * _J * _dt

	mpo1 = MPO_1(_N_qubit, _theta1,
				 _upper_ind_id=_upper_ind_id, _lower_ind_id="_tem_c{}", _site_tag_id=_site_tag_id)
	mpo2 = MPO_2(_N_qubit, _theta2,
				 _upper_ind_id=_upper_ind_id, _lower_ind_id="_tem_c{}", _site_tag_id=_site_tag_id)
	mpo3 = MPO_3(_N_qubit,
				 _upper_ind_id=_upper_ind_id, _lower_ind_id="_tem_c{}", _site_tag_id=_site_tag_id)
	mpo4 = MPO_4(_N_qubit, _theta2,
				 _upper_ind_id=_upper_ind_id, _lower_ind_id="_tem_c{}", _site_tag_id=_site_tag_id)

	inv_mpo1 = inv_MPO_1(_N_qubit, _theta1,
						 _upper_ind_id="_tem_b{}", _lower_ind_id="_tem_a{}", _site_tag_id=_site_tag_id)
	inv_mpo2 = inv_MPO_2(_N_qubit, _theta2,
						 _upper_ind_id="_tem_b{}", _lower_ind_id="_tem_a{}", _site_tag_id=_site_tag_id)
	inv_mpo3 = MPO_3(_N_qubit,
					 _upper_ind_id="_tem_b{}", _lower_ind_id="_tem_a{}", _site_tag_id=_site_tag_id)
	inv_mpo4 = inv_MPO_4(_N_qubit, _theta2,
						 _upper_ind_id="_tem_b{}", _lower_ind_id="_tem_a{}", _site_tag_id=_site_tag_id)
	inv_splm1 = inv_SPLM_ptm_mpo(_N_qubit, _layer1_lam1_mat, _layer1_lam2_mat,
								 upper_ind_id="_tem_a{}", lower_ind_id=_lower_ind_id, site_tag_id=_site_tag_id)
	inv_splm2 = inv_SPLM_ptm_mpo(_N_qubit, _layer2_lam1_mat, _layer2_lam2_mat,
								 upper_ind_id="_tem_a{}", lower_ind_id=_lower_ind_id, site_tag_id=_site_tag_id)

	index_map1 = {_old_upper_ind: _new_upper_ind for _old_upper_ind, _new_upper_ind in
				  map(lambda _: ("_tem_a{}".format(_), _upper_ind_id.format(_)), range(_N_qubit))}
	index_map2 = {_old_upper_ind: _new_upper_ind for _old_upper_ind, _new_upper_ind in
				  map(lambda _: (_upper_ind_id.format(_), "_tem_c{}".format(_)), range(_N_qubit))} \
				 | {_old_lower_ind: _new_lower_ind for _old_lower_ind, _new_lower_ind in
					map(lambda _: (_lower_ind_id.format(_), "_tem_b{}".format(_)), range(_N_qubit))}
	_site_tags = tuple(map(_site_tag_id.format, range(_N_qubit)))

	def TEM_MPO(_p_tn):
		if _p_tn is None:
			_tn = inv_splm1.reindex(index_map1)
		else:
			_tn = inv_splm1 | inv_mpo1 | _p_tn.reindex(index_map2) | mpo1
			qtn.tensor_network_1d_compress(_tn, max_bond=_max_bond, cutoff=_cutoff, site_tags=_site_tags, method=method,
										   inplace=True)

		_tn.reindex(index_map2, inplace=True)
		_tn = inv_splm1 | inv_mpo2 | _tn | mpo2
		qtn.tensor_network_1d_compress(_tn, max_bond=_max_bond, cutoff=_cutoff, site_tags=_site_tags, method=method,
									   inplace=True)

		_tn.reindex(index_map2, inplace=True)
		_tn = inv_splm2 | inv_mpo3 | _tn | mpo3
		qtn.tensor_network_1d_compress(_tn, max_bond=_max_bond, cutoff=_cutoff, site_tags=_site_tags, method=method,
									   inplace=True)

		_tn.reindex(index_map2, inplace=True)
		_tn = inv_splm2 | inv_mpo4 | _tn | mpo4
		qtn.tensor_network_1d_compress(_tn, max_bond=_max_bond, cutoff=_cutoff, site_tags=_site_tags, method=method,
									   inplace=True)
		return _tn

	return TEM_MPO

def TEM_trotter_step_MPO(
		_N_qubit, _J: float, _h: float, _dt: float,
		_layer1_lam1_mat, _layer1_lam2_mat, _layer2_lam1_mat, _layer2_lam2_mat,
		_left_upper_ind_id="kl{}", _left_lower_ind_id="bl{}",
		_right_upper_ind_id="kr{}", _right_lower_ind_id="br{}",
		_site_tag_id="I{}",
		_max_bond: int = 200, _cutoff: float = 1e-10, method='direct'
):
	from SPLM import inv_SPLM_ptm_mpo

	_theta1, _theta2 = 2 * _h * _dt, -2 * _J * _dt
	_site_tags = tuple(map(_site_tag_id.format, range(_N_qubit)))

	mpo_l = inv_SPLM_ptm_mpo(_N_qubit, _layer2_lam1_mat, _layer2_lam2_mat,
							 upper_ind_id="_tem_a{}", lower_ind_id=_left_lower_ind_id, site_tag_id=_site_tag_id) \
			| inv_MPO_4(_N_qubit, _theta2,
						_upper_ind_id="_tem_b{}", _lower_ind_id="_tem_a{}", _site_tag_id=_site_tag_id) \
			| inv_SPLM_ptm_mpo(_N_qubit, _layer2_lam1_mat, _layer2_lam2_mat,
							   upper_ind_id="_tem_c{}", lower_ind_id="_tem_b{}", site_tag_id=_site_tag_id) \
			| MPO_3(_N_qubit,
					_upper_ind_id="_tem_d{}", _lower_ind_id="_tem_c{}", _site_tag_id=_site_tag_id) \
			| inv_SPLM_ptm_mpo(_N_qubit, _layer1_lam1_mat, _layer1_lam2_mat,
							   upper_ind_id="_tem_e{}", lower_ind_id="_tem_d{}", site_tag_id=_site_tag_id) \
			| inv_MPO_2(_N_qubit, _theta2,
						_upper_ind_id="_tem_f{}", _lower_ind_id="_tem_e{}", _site_tag_id=_site_tag_id) \
			| inv_SPLM_ptm_mpo(_N_qubit, _layer1_lam1_mat, _layer1_lam2_mat,
							   upper_ind_id="_tem_g{}", lower_ind_id="_tem_f{}", site_tag_id=_site_tag_id) \
			| inv_MPO_1(_N_qubit, _theta1,
						_upper_ind_id=_left_upper_ind_id, _lower_ind_id="_tem_g{}", _site_tag_id=_site_tag_id)
	qtn.tensor_network_1d_compress(mpo_l, max_bond=_max_bond, cutoff=_cutoff, site_tags=_site_tags, method=method,
								   inplace=True)

	mpo_r = MPO_1(_N_qubit, _theta1,
				  _upper_ind_id="_tem_h{}", _lower_ind_id=_right_lower_ind_id, _site_tag_id=_site_tag_id) \
			| MPO_2(_N_qubit, _theta2,
					_upper_ind_id="_tem_i{}", _lower_ind_id="_tem_h{}", _site_tag_id=_site_tag_id) \
			| MPO_3(_N_qubit,
					_upper_ind_id="_tem_j{}", _lower_ind_id="_tem_i{}", _site_tag_id=_site_tag_id) \
			| MPO_4(_N_qubit, _theta2,
					_upper_ind_id=_right_upper_ind_id, _lower_ind_id="_tem_j{}", _site_tag_id=_site_tag_id)
	qtn.tensor_network_1d_compress(mpo_r, max_bond=_max_bond, cutoff=_cutoff, site_tags=_site_tags, method=method,
								   inplace=True)

	# index_map1 = {_old_upper_ind: _new_upper_ind for _old_upper_ind, _new_upper_ind in
	# 			  map(lambda _: ("_tem_a{}".format(_), _upper_ind_id.format(_)), range(_N_qubit))}
	# index_map2 = {_old_upper_ind: _new_upper_ind for _old_upper_ind, _new_upper_ind in
	# 			  map(lambda _: (_upper_ind_id.format(_), "_tem_c{}".format(_)), range(_N_qubit))} \
	# 			 | {_old_lower_ind: _new_lower_ind for _old_lower_ind, _new_lower_ind in
	# 				map(lambda _: (_lower_ind_id.format(_), "_tem_b{}".format(_)), range(_N_qubit))}
	return mpo_l, mpo_r
